package tokenBuket

import (
	"context"
	"fmt"
	"github.com/go-redis/redis/v8"
	"queue-gateway/common"
	"queue-gateway/config"
	"queue-gateway/databases"
	"strconv"
	"sync"
	"time"
)

var serviceConfig = config.ServiceConfig
var prefix = config.App.Project

var ServiceMap = sync.Map{}

type Service struct {
	name       string
	addressMap map[string]struct{}
	HashName   string // redis 令牌桶中对应的hash key
	limit      int    // token limit
	timeout    int
}

func UpdateServiceMap(ctx context.Context, key string, PodIps []string, resourceVersion string) {
	zsetName := fmt.Sprintf("%s_%s", prefix, key)
	limit := 0
	timeout := 60

	for _, v := range serviceConfig.Services {
		if v.Name == key {
			limit = v.TokenCapacity
			timeout = v.TimeOut

		}
	}
	podIpMap := make(map[string]struct{})
	for _, v2 := range PodIps {
		podIpMap[v2] = struct{}{}
	}
	ServiceMap.Store(key, &Service{name: key, addressMap: podIpMap, HashName: zsetName, limit: limit, timeout: timeout})
	err := updateServiceHashAndServiceTB(ctx, key, resourceVersion)
	if err != nil {
		return
	}
}

func updateServiceHashAndServiceTB(ctx context.Context, svr, resourceVersion string) error {
	serTemp, _ := ServiceMap.Load(svr)
	service := serTemp.(*Service)
	hashName := service.HashName
	err := AcquireLock(ctx, resourceVersion)
	time.Sleep(time.Second * 2)
	if err != nil {
		return err
	}
	hm, err := databases.Rdb.ZRange(ctx, hashName, 0, -1).Result()
	if err != nil {
		return err
	}
	for podip, _ := range service.addressMap {
		script := "local v = redis.call('zrank', KEYS[1], ARGV[1]) if not v then " +
			"redis.call('zadd', KEYS[1], 0, ARGV[1]) end"
		databases.Rdb.Eval(ctx, script, []string{hashName}, podip)
	}

	// 删除redis里面对应的service hash中多余的podip key
	for _, oldPodip := range hm {
		if _, ok := service.addressMap[oldPodip]; !ok {
			script := `local v = redis.call('zrank', KEYS[1], ARGV[1]);
	if v then
	   redis.call('zrem', KEYS[1], ARGV[1])
	   redis.call('del', KEYS[2])
	   return 0
	end return 1`
			podipKey := fmt.Sprintf("%s_%s", prefix, oldPodip)
			databases.Rdb.Eval(ctx, script, []string{hashName, podipKey}, oldPodip)
		}
	}
	//time.Sleep(time.Second * 15)
	ReleaseLock(ctx, resourceVersion)
	return nil
}

type TokenBuket struct {
	upstream string
	services map[string]*Service
}

func popToken(bucket chan struct{}, n int) {
	for i := 0; i < n; i++ {
		<-bucket
	}

}

func RevertToken(ctx context.Context, zsetName, podIp string, tokenId interface{}) error {
	_, err := databases.Rdb.TxPipelined(ctx, func(pipe redis.Pipeliner) error {
		pipe.ZIncrBy(ctx, zsetName, -1, podIp)
		podipKey := fmt.Sprintf("%s_%s", prefix, podIp)
		pipe.LRem(ctx, podipKey, 1, tokenId)
		return nil
	})
	if err != nil {
		return err
	}
	return nil
}

func GetToken(ctx context.Context, services []string, weights []float32) (string, string, int64) {
	serviceIndex := common.WeightedRandomIndex(weights)
	svr := services[serviceIndex]
	s, ok := ServiceMap.Load(svr)
	if !ok {
		return "", "", 0
	}
	service, ok := s.(*Service)
	podip, tokenId := GetPodIp(ctx, service.HashName, service.limit)
	fmt.Println("获取到的podip是", podip)
	return service.HashName, podip, tokenId
}

func GetPodIp(ctx context.Context, hashName string, limit int) (string, int64) {
	tokenId := time.Now().Unix()
	script := "local e= redis.call('zrange', KEYS[1], 0, 0, 'withscores') if #e==2 then if e[2]<ARGV[1] then local " +
		"lp = e[1] local t = {} table.insert(t, KEYS[2]) table.insert(t, lp) local newlp = table.concat(t) " +
		"redis.call('rpush', newlp, ARGV[2]) redis.call('zincrby', KEYS[1], 1, e[1]) return e[1] end end return ''"
	podipPrefix := fmt.Sprintf("%s_", prefix)
	cmd := databases.Rdb.Eval(ctx, script, []string{hashName, podipPrefix}, limit, tokenId)
	result := cmd.Val()
	podip := result.(string)
	return podip, tokenId
}

// ReleaseToken 异步线程清除令牌桶中过期占用的令牌
func ReleaseToken(ctx context.Context) {
LOOP:
	for {
		select {
		case <-ctx.Done():
			break LOOP
		default:
			ServiceMap.Range(func(key, value any) bool {
				service := value.(*Service)
				for ip, _ := range service.addressMap {
					podipKey := fmt.Sprintf("%s_%s", prefix, ip)
					cmd := databases.Rdb.LRange(ctx, podipKey, 0, -1)
					if len(cmd.Val()) > 0 {
						for _, v := range cmd.Val() {
							nowts := time.Now().Unix()
							ts, _ := strconv.Atoi(v)
							if (int(nowts) - ts) > service.timeout {
								databases.Rdb.TxPipelined(ctx, func(pipe redis.Pipeliner) error {
									pipe.ZIncrBy(ctx, service.HashName, -1, ip)
									pipe.LRem(ctx, podipKey, 1, v)
									return nil
								})

							}
						}
					}
				}
				return true
			})
			time.Sleep(time.Second)
		}
	}
}
